# -*- coding: utf-8 -*-
"""
A股股票信息同步模块
从akshare或tushare获取A股股票基本信息，并同步到new_stock_info表
"""
import logging
import pandas as pd
import akshare as ak
import tushare as ts
from datetime import datetime
from typing import Dict, Any
from database import DatabaseManager, StockInfoDAO
from config import DATABASE_CONFIG, SYNC_CONFIG, DATA_SOURCE_CONFIG
import time

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class StockInfoSync:
    """A股股票信息同步器"""
    
    def __init__(self):
        self.db_manager = DatabaseManager(DATABASE_CONFIG)
        self.stock_dao = StockInfoDAO(self.db_manager)
        self.batch_size = SYNC_CONFIG.get('batch_size', 1000)
        
        # 数据源配置
        self.data_source = DATA_SOURCE_CONFIG.get('data_source', 'tushare')
        self.tushare_token = DATA_SOURCE_CONFIG.get('tushare_token', '')
        self.request_delay = DATA_SOURCE_CONFIG.get('request_delay', 0.35)
        
        # 初始化 Tushare
        if self.data_source == 'tushare' and self.tushare_token:
            try:
                ts.set_token(self.tushare_token)
                self.tushare_pro = ts.pro_api()
                logger.info("Tushare API 初始化成功")
            except Exception as e:
                logger.error(f"Tushare API 初始化失败: {e}")
                logger.warning("将回退到 akshare 数据源")
                self.data_source = 'akshare'
        elif self.data_source == 'tushare' and not self.tushare_token:
            logger.warning("未配置 Tushare token，将使用 akshare 数据源")
            self.data_source = 'akshare'
        
        logger.info(f"当前使用数据源: {self.data_source}")
    
    def get_all_stocks_tushare(self) -> pd.DataFrame:
        """使用Tushare获取所有A股信息
        
        注意：Tushare的stock_basic接口限制为每小时访问1次（免费用户）
        建议：股票列表不需要频繁更新，每天或每周同步一次即可
        """
        try:
            logger.info("正在使用Tushare获取所有A股信息...")
            logger.info("提示：stock_basic接口限制为每小时1次，如触发限制将自动切换到AKShare")
            
            # 获取股票列表基本信息
            # list_status: L=上市 D=退市 P=暂停上市
            df = self.tushare_pro.stock_basic(
                exchange='',
                list_status='L',
                fields='ts_code,symbol,name,area,industry,market,list_date'
            )
            
            time.sleep(self.request_delay)
            
            if df is None or df.empty:
                logger.warning("Tushare 返回空数据")
                return pd.DataFrame()
            
            # 字段映射
            df.rename(columns={
                'symbol': 'stock_code',
                'name': 'stock_name',
                'list_date': 'listing_date'
            }, inplace=True)
            
            # 根据 ts_code 判断交易所
            def get_exchange(ts_code):
                if ts_code.endswith('.SH'):
                    return 'SHSE'
                elif ts_code.endswith('.SZ'):
                    return 'SZSE'
                elif ts_code.endswith('.BJ'):
                    return 'BSE'  # 北交所
                else:
                    return None
            
            df['exchange'] = df['ts_code'].apply(get_exchange)
            
            # 根据股票代码判断板块
            def get_market_type(code):
                code_str = str(code).zfill(6)
                if code_str.startswith('000'):
                    return "主板"
                elif code_str.startswith('002'):
                    return "中小板"
                elif code_str.startswith('300'):
                    return "创业板"
                elif code_str.startswith('600') or code_str.startswith('601') or code_str.startswith('603') or code_str.startswith('605'):
                    return "主板"
                elif code_str.startswith('688'):
                    return "科创板"
                elif code_str.startswith('8') or code_str.startswith('4'):
                    return "北交所"
                else:
                    return "未知"
            
            df['market'] = df['stock_code'].apply(get_market_type)
            
            # 只保留上交所、深交所的股票，过滤掉北交所等
            filtered_df = df[df['exchange'].isin(['SHSE', 'SZSE'])].copy()
            
            # 处理上市日期格式 (YYYYMMDD -> YYYY-MM-DD)
            def format_listing_date(date_str):
                if pd.isna(date_str) or date_str == '':
                    return None
                try:
                    date_str = str(date_str)
                    return f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:8]}"
                except:
                    return None
            
            filtered_df['listing_date'] = filtered_df['listing_date'].apply(format_listing_date)
            
            logger.info(f"Tushare 获取原始A股信息: {len(df)} 只股票")
            logger.info(f"过滤后保留股票: {len(filtered_df)} 只股票")
            logger.info(f"过滤掉的股票: {len(df) - len(filtered_df)} 只")
            
            return filtered_df
            
        except Exception as e:
            error_msg = str(e)
            logger.error(f"Tushare获取所有A股信息失败: {e}")
            
            # 检查是否是频率限制错误
            if "每小时最多访问" in error_msg or "每分钟最多访问" in error_msg:
                logger.warning("触发Tushare频率限制，自动切换到AKShare数据源")
                return self.get_all_stocks_akshare()
            
            raise
    
    def get_all_stocks_akshare(self) -> pd.DataFrame:
        """使用AKShare获取所有A股信息（只包括上交所、深交所的主板、科创板、创业板）"""
        try:
            logger.info("正在使用AKShare获取所有A股信息...")
            all_df = ak.stock_info_a_code_name()
            
            # 重命名列
            all_df.rename(columns={
                "code": "stock_code", 
                "name": "stock_name"
            }, inplace=True)
            
            # 根据股票代码判断交易所和板块，并过滤不需要的股票
            def get_exchange_and_market(code):
                code_str = str(code).zfill(6)
                if code_str.startswith('000'):  # 深交所主板
                    return "SZSE", "主板"
                elif code_str.startswith('002'):  # 深交所中小板
                    return "SZSE", "中小板"
                elif code_str.startswith('300'):  # 深交所创业板
                    return "SZSE", "创业板"
                elif code_str.startswith('600') or code_str.startswith('601') or code_str.startswith('603') or code_str.startswith('605'):  # 上交所主板
                    return "SHSE", "主板"
                elif code_str.startswith('688'):  # 上交所科创板
                    return "SHSE", "科创板"
                else:
                    # 过滤掉新三板、北交所等不需要的股票
                    return None, None
            
            # 应用函数并过滤
            exchange_market = all_df['stock_code'].apply(get_exchange_and_market)
            all_df['exchange'] = [x[0] for x in exchange_market]
            all_df['market'] = [x[1] for x in exchange_market]
            
            # 只保留上交所、深交所的股票
            filtered_df = all_df[all_df['exchange'].notna()].copy()
            
            # 添加其他字段
            filtered_df['industry'] = None
            filtered_df['listing_date'] = None
            
            logger.info(f"AKShare 获取原始A股信息: {len(all_df)} 只股票")
            logger.info(f"过滤后保留股票: {len(filtered_df)} 只股票")
            logger.info(f"过滤掉的股票: {len(all_df) - len(filtered_df)} 只")
            
            return filtered_df
            
        except Exception as e:
            logger.error(f"AKShare获取所有A股信息失败: {e}")
            raise
    
    def get_all_stocks(self) -> pd.DataFrame:
        """获取所有A股信息（根据配置选择数据源）"""
        if self.data_source == 'tushare':
            return self.get_all_stocks_tushare()
        else:
            return self.get_all_stocks_akshare()
    
    def _get_sz_market(self, code: str) -> str:
        """根据深交所股票代码判断板块"""
        if code.startswith('000'):
            return "主板"
        elif code.startswith('002'):
            return "中小板"
        elif code.startswith('300'):
            return "创业板"
        else:
            return "未知"
    
    def process_stock_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """处理股票数据，补充必要字段并识别退市股票"""
        try:
            # 复制数据框
            processed_df = df.copy()
            
            # 确保股票代码格式正确（只保留6位数字）
            def format_stock_code(row):
                code = str(row['stock_code']).zfill(6)  # 补齐6位数字
                return code
            
            processed_df['stock_code'] = processed_df.apply(format_stock_code, axis=1)
            
            # 识别退市股票状态
            def determine_stock_status(row):
                stock_name = str(row['stock_name'])
                # 检查是否包含退市相关标识
                if any(keyword in stock_name for keyword in ['*ST', 'ST', 'st', '退市', '暂停']):
                    return "delisted"
                else:
                    return "active"
            
            processed_df["status"] = processed_df.apply(determine_stock_status, axis=1)
            processed_df["updated_at"] = datetime.now()
            
            # 处理日期格式
            if 'listing_date' in processed_df.columns:
                # 将listing_date设置为None，因为ak.stock_info_a_code_name()不包含上市日期
                processed_df['listing_date'] = None
            
            # 数据清洗
            processed_df = processed_df.dropna(subset=['stock_code', 'stock_name'])
            processed_df = processed_df.drop_duplicates(subset=['stock_code'])
            
            # 统计退市股票数量
            delisted_count = len(processed_df[processed_df['status'] == 'delisted'])
            active_count = len(processed_df[processed_df['status'] == 'active'])
            
            # 选择需要的列
            required_columns = ['stock_code', 'stock_name', 'exchange', 'industry', 'market', 'listing_date', 'status', 'updated_at']
            processed_df = processed_df[required_columns]
            
            logger.info(f"数据处理完成，共 {len(processed_df)} 条有效记录")
            logger.info(f"正常股票: {active_count} 只，退市/ST股票: {delisted_count} 只")
            return processed_df
            
        except Exception as e:
            logger.error(f"处理股票数据失败: {e}")
            raise
    
    def check_and_cleanup_unwanted_stocks(self) -> int:
        """检查并清理数据库中不需要的股票（新三板、北交所等）"""
        try:
            # 先检查是否存在不需要的股票
            check_sql = """
            SELECT COUNT(*) as count 
            FROM new_stock_info 
            WHERE stock_code LIKE '8%' OR stock_code LIKE '4%' OR stock_code LIKE '9%'
            """
            
            unwanted_count = self.db_manager.execute_query(check_sql)[0]['count']
            
            if unwanted_count == 0:
                logger.info("数据库中无需要清理的股票，跳过清理步骤")
                return 0
            
            logger.info(f"发现 {unwanted_count} 只需要清理的股票，开始清理...")
            
            # 删除新三板、北交所等不需要的股票
            cleanup_sql = """
            DELETE FROM new_stock_info 
            WHERE stock_code LIKE '8%' OR stock_code LIKE '4%' OR stock_code LIKE '9%'
            """
            
            deleted_count = self.db_manager.execute_update(cleanup_sql)
            logger.info(f"已清理 {deleted_count} 只不需要的股票")
            
            return deleted_count
            
        except Exception as e:
            logger.error(f"检查并清理不需要的股票失败: {e}")
            return 0
    
    def mark_delisted_stocks(self) -> int:
        """标记数据库中已退市但仍在数据库中的股票"""
        try:
            logger.info("开始标记退市股票...")
            
            # 获取当前akshare中的所有股票代码
            current_stocks = self.get_all_stocks()
            current_codes = set(current_stocks['stock_code'].astype(str))
            
            # 获取数据库中的所有股票代码
            db_codes = set(self.stock_dao.get_all_stock_codes())
            
            # 找出在数据库中但不在akshare中的股票（可能已退市）
            potentially_delisted = db_codes - current_codes
            
            if potentially_delisted:
                # 将这些股票标记为退市
                delisted_codes = list(potentially_delisted)
                placeholders = ','.join(['%s'] * len(delisted_codes))
                update_sql = f"""
                UPDATE new_stock_info 
                SET status = 'delisted', updated_at = CURRENT_TIMESTAMP
                WHERE stock_code IN ({placeholders})
                """
                
                updated_count = self.db_manager.execute_update(update_sql, tuple(delisted_codes))
                logger.info(f"已标记 {updated_count} 只股票为退市状态")
                return updated_count
            else:
                logger.info("没有发现新的退市股票")
                return 0
                
        except Exception as e:
            logger.error(f"标记退市股票失败: {e}")
            return 0
    
    def sync_stock_info(self) -> Dict[str, Any]:
        """同步A股股票信息到new_stock_info表"""
        start_time = datetime.now()
        result = {
            'success': False,
            'total_count': 0,
            'sz_count': 0,
            'sh_count': 0,
            'synced_count': 0,
            'cleaned_count': 0,
            'delisted_count': 0,
            'error': None,
            'duration': 0
        }
        
        try:
            logger.info("开始同步A股股票信息...")
            
            # 1. 检查并清理不需要的股票
            cleaned_count = self.check_and_cleanup_unwanted_stocks()
            result['cleaned_count'] = cleaned_count
            
            # 2. 获取所有A股信息
            all_df = self.get_all_stocks()
            result['total_count'] = len(all_df)
            
            # 统计各交易所数量
            result['sz_count'] = len(all_df[all_df['exchange'] == 'SZSE'])
            result['sh_count'] = len(all_df[all_df['exchange'] == 'SHSE'])
            
            # 3. 处理数据
            processed_df = self.process_stock_data(all_df)
            
            # 4. 批量同步到数据库
            logger.info(f"开始批量同步 {len(processed_df)} 条股票信息到数据库...")
            
            # 分批处理
            total_synced = 0
            for i in range(0, len(processed_df), self.batch_size):
                batch_df = processed_df.iloc[i:i + self.batch_size]
                batch_data = batch_df.to_dict('records')
                
                try:
                    synced_count = self.stock_dao.batch_insert_or_update_stock_info(batch_data)
                    total_synced += synced_count
                    logger.info(f"已同步第 {i//self.batch_size + 1} 批，共 {synced_count} 条记录")
                    
                except Exception as e:
                    logger.error(f"同步第 {i//self.batch_size + 1} 批失败: {e}")
                    continue
            
            # 5. 标记退市股票
            delisted_count = self.mark_delisted_stocks()
            result['delisted_count'] = delisted_count
            
            result['success'] = True
            result['synced_count'] = total_synced
            logger.info(f"A股股票信息同步完成！")
            logger.info(f"同步记录: {total_synced} 条")
            logger.info(f"清理记录: {cleaned_count} 条")
            logger.info(f"退市标记: {delisted_count} 条")
            
        except Exception as e:
            result['error'] = str(e)
            logger.error(f"A股股票信息同步失败: {e}")
        
        finally:
            end_time = datetime.now()
            result['duration'] = (end_time - start_time).total_seconds()
            logger.info(f"同步耗时: {result['duration']:.2f} 秒")
        
        return result
    
    def get_sync_summary(self) -> Dict[str, Any]:
        """获取同步摘要信息"""
        try:
            # 获取数据库中的股票总数和按交易所统计
            total_count = self.db_manager.execute_query('SELECT COUNT(*) as count FROM new_stock_info WHERE status = "active"')[0]['count']
            sz_count = self.db_manager.execute_query('SELECT COUNT(*) as count FROM new_stock_info WHERE status = "active" AND exchange = "SZSE"')[0]['count']
            sh_count = self.db_manager.execute_query('SELECT COUNT(*) as count FROM new_stock_info WHERE status = "active" AND exchange = "SHSE"')[0]['count']
            
            return {
                'total_stocks': total_count,
                'sz_stocks': sz_count,
                'sh_stocks': sh_count,
                'last_sync': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            
        except Exception as e:
            logger.error(f"获取同步摘要失败: {e}")
            return {
                'error': str(e),
                'total_stocks': 0,
                'sz_stocks': 0,
                'sh_stocks': 0
            }

def main():
    """主函数"""
    try:
        # 创建同步器
        sync = StockInfoSync()
        
        # 执行同步
        result = sync.sync_stock_info()
        
        # 打印结果
        if result['success']:
            print("=" * 50)
            print("A股股票信息同步成功！")
            print(f"深交所股票: {result['sz_count']} 只")
            print(f"上交所股票: {result['sh_count']} 只")
            print(f"总计: {result['total_count']} 只")
            print(f"同步记录: {result['synced_count']} 条")
            print(f"清理记录: {result['cleaned_count']} 条")
            print(f"退市标记: {result['delisted_count']} 条")
            print(f"同步耗时: {result['duration']:.2f} 秒")
            print("=" * 50)
        else:
            print("=" * 50)
            print("A股股票信息同步失败！")
            print(f"错误信息: {result['error']}")
            print("=" * 50)
        
        # 显示同步摘要
        summary = sync.get_sync_summary()
        if 'error' not in summary:
            print("\n数据库当前状态:")
            print(f"总股票数: {summary['total_stocks']}")
            print(f"深交所: {summary['sz_stocks']} 只")
            print(f"上交所: {summary['sh_stocks']} 只")
        
    except Exception as e:
        logger.error(f"程序执行失败: {e}")
        print(f"程序执行失败: {e}")

if __name__ == "__main__":
    main()
